"""
1.	YOLOv1网络结构实现。YOLOv1的前20层是用于特征提取的，也就是随便替换为一个分类网络(除去最后的全连接层)其实都行。
这里使用ResNet34的网络作为特征提取部分。这样做的好处是，pytorch的torchvision中提供了ResNet34的预训练模型。然后，
除去ResNet34的最后两层，再连接上YOLOv1的最后4个卷积层和两个全连接层，作为YOLOv1网络结构。具体实现代码如下：
"""
import torch as pt
import torchvision as ptv

# ①	定义主程序入口
if '__main__' == __name__:

    # ②	定义类函数YOLOv1_resnet（）
    class YOLOv1_resnet(pt.nn.Module):

        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            # ③	调用torchvision里的resnet34预训练模型
            resnet = ptv.models.resnet34(pretrained=True)
            for p in resnet.parameters():
                p.requires_grad = False
            # ④	记录resnet全连接层之前的网络输出通道数，方便连入后续卷积网络中
            resnet_out_ch = resnet.fc.in_features
            print('resnet_out_ch', resnet_out_ch)
            # ⑤	去除resnet的最后两层
            self.resnet = pt.nn.Sequential(*(list(resnet.children())[:-2]))
            # ⑥	用nn.Sequential()，定义YOLOv1的最后4个卷积层
            self.customer_cnn = pt.nn.Sequential(
                pt.nn.Conv2d(resnet_out_ch, 512, (3, 3), (1, 1), 1),
                pt.nn.BatchNorm2d(512),
                pt.nn.ReLU(),
                pt.nn.Conv2d(512, 512, (3, 3), (2, 2), 1),
                pt.nn.BatchNorm2d(512),
                pt.nn.ReLU(),
                pt.nn.Conv2d(512, 1024, (3, 3), (1, 1), 1),
                pt.nn.BatchNorm2d(1024),
                pt.nn.ReLU(),
                pt.nn.Conv2d(1024, 1024, (3, 3), (1, 1), 1),
                pt.nn.BatchNorm2d(1024),
                pt.nn.ReLU(),
            )
            # ⑦	用nn.Sequential()，定义YOLOv1的最后2个全连接层
            self.customer_fc = pt.nn.Sequential(
                pt.nn.Linear(50176, 1024),
                pt.nn.ReLU(),
                pt.nn.Linear(1024, 1470),
                pt.nn.Sigmoid(),
            )

        # ⑧	进行前向传播
        def forward(self, x):
            x = self.resnet(x)
            print('after resnet', x.size())
            x = self.customer_cnn(x)
            x = x.reshape(-1, 50176)
            x = self.customer_fc(x)
            return x


    # ⑨	自定义输入张量（(5,3,448,448)，验证网络可以正常跑通
    model = YOLOv1_resnet()
    x = pt.zeros((5, 3, 448, 448), dtype=pt.float32)
    pred = model(x)

    # ⑩	在前向传播中，打印resnet34输出维度和全连接输出维度
    print('after fc', pred.size())
